import os
import cv2
import numpy as np
import torch
from .net.dehaze_net import dehaze_net  # 正确：导入类 `dehaze_net`
from PIL import Image  # 导入 Image 模块
def process_image(input_path, output_dir):
    # 加载模型
    model = dehaze_net().cuda()
    # 构造模型权重文件的绝对路径
    model_path = os.path.join(os.path.dirname(__file__), 'snapshots', 'dehazer.pth')
    model.load_state_dict(torch.load(model_path))
    model.eval()

    # 读取图像
    img = Image.open(input_path)
    img = np.asarray(img) / 255.0
    img = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0).cuda()

    # 处理图像
    with torch.no_grad():
        clean_image = model(img)

    # 保存处理后的图像
    clean_image = clean_image.cpu().squeeze(0).permute(1, 2, 0).numpy()
    clean_image = (clean_image * 255).astype(np.uint8)

    filename = os.path.basename(input_path)
    output_filename = f"processed_{filename}"
    output_path = os.path.join(output_dir, output_filename)
    cv2.imwrite(output_path, cv2.cvtColor(clean_image, cv2.COLOR_RGB2BGR))

    return output_filename